//
//  Untitled.swift
//  PLDownload
//
//  Created by bin hu on 2025/7/7.
//

/**
 流媒体处理
 1.解析m3u8文件
 2.获取TS片段列表
 3.创建分段下载任务
 4.AES解密（ if need ）
 5.合并TS文件
 */
//class M3U8Handler: ProtocolHandler {
//    private let keyManager = KeyManager.shared
//    private let parser = M3U8Parser()
//    private let decryptor = AESDecryptor()
//    private var activeSegmentDownloads = [URLSessionTask]()
//    
//    func startDownload(task: DownloadTask) {
//        // 1. 下载m3u8索引文件
//        downloadIndexFile(task: task)
//    }
//    
//    private func downloadIndexFile(task: DownloadTask) {
//        let indexTask = URLSession.shared.dataTask(with: task.url) { [weak self] data, _, error in
//            guard let self = self, let data = data,
//                  let indexContent = String(data: data, encoding: .utf8) else {
//                task.completion(.failure(error ?? DownloadError.m3u8ParseFailed))
//                return
//            }
//            
//            // 2. 解析m3u8内容
//            let encryptedM3U8 = self.parser.parseEncryptedM3U8(
//                content: indexContent,
//                baseURL: task.url.deletingLastPathComponent()
//            )
//            
//            // 3. 处理加密内容
//            self.handleEncryption(for: task, encryptedM3U8: encryptedM3U8)
//        }
//        indexTask.resume()
//    }
//    
//    private func handleEncryption(for task: DownloadTask, encryptedM3U8: EncryptedM3U8) {
//        guard let method = encryptedM3U8.method, method == "AES-128" else {
//            // 非加密或不受支持的加密方式
//            downloadSegments(task: task, segments: encryptedM3U8.tsSegments)
//            return
//        }
//        
//        guard let keyURI = encryptedM3U8.keyURI else {
//            task.completion(.failure(DownloadError.missingEncryptionKey))
//            return
//        }
//        
//        // 4. 获取解密密钥
//        keyManager.getKey(for: keyURI) { [weak self] result in
//            switch result {
//            case .success(let keyData):
//                // 获取IV（初始向量）
//                let iv = encryptedM3U8.iv ?? self?.generateDefaultIV(for: keyData) ?? Data()
//                
//                // 5. 下载并解密片段
//                self?.downloadAndDecryptSegments(
//                    task: task,
//                    segments: encryptedM3U8.tsSegments,
//                    key: keyData,
//                    iv: iv
//                )
//                
//            case .failure(let error):
//                task.completion(.failure(error))
//            }
//        }
//    }
//    
//    private func generateDefaultIV(for key: Data) -> Data {
//        // 根据HLS规范，当IV未指定时，使用片段的序列号
//        // 这里简化处理：使用密钥的MD5作为默认IV
//        return key.md5()
//    }
//}
//
///// 动态iv处理
//extension M3U8Handler {
//    private func downloadAndDecryptSegments(task: DownloadTask,
//                                          segments: [URL],
//                                          key: Data,
//                                          iv: Data) {
//        let group = DispatchGroup()
//        let tempDir = createTempDirectory(for: task)
//        var segmentPaths = [URL]()
//        
//        for (index, segmentURL) in segments.enumerated() {
//            group.enter()
//            
//            // 为每个片段生成唯一IV（HLS规范要求）
//            let segmentIV = generateSegmentIV(baseIV: iv, sequence: index)
//            
//            let segmentTask = URLSession.shared.dataTask(with: segmentURL) { [weak self] data, _, error in
//                defer { group.leave() }
//                
//                guard let self = self, let data = data, error == nil else {
//                    // 错误处理
//                    return
//                }
//                
//                do {
//                    // 解密TS片段
//                    let decryptedData = try self.decryptor.decryptTS(
//                        data: data,
//                        key: key,
//                        iv: segmentIV
//                    )
//                    
//                    // 验证解密结果
//                    guard self.decryptor.validateDecryption(decryptedData: decryptedData) else {
//                        throw DownloadError.decryptionFailed
//                    }
//                    
//                    // 保存解密后的片段
//                    let segmentPath = tempDir.appendingPathComponent("seg_\(index).ts")
//                    try decryptedData.write(to: segmentPath)
//                    segmentPaths.append(segmentPath)
//                    
//                } catch {
//                    // 错误处理
//                }
//            }
//            segmentTask.resume()
//            activeSegmentDownloads.append(segmentTask)
//        }
//        
//        group.notify(queue: .main) {
//            // 6. 合并所有TS文件
//            self.mergeSegments(task: task, segmentPaths: segmentPaths)
//        }
//    }
//    
//    private func generateSegmentIV(baseIV: Data, sequence: Int) -> Data {
//        // 根据HLS规范生成片段特定IV
//        var mutableIV = baseIV
//        for i in 0..<MemoryLayout<Int>.size {
//            mutableIV[i] ^= UInt8((sequence >> (i * 8)) & 0xFF)
//        }
//        return mutableIV
//    }
//}
//
///// 分片下载与合并
//extension M3U8Handler {
//    private func mergeSegments(task: DownloadTask, segmentPaths: [URL]) {
//        let outputPath = task.finalPath
//        
//        do {
//            let fileManager = FileManager.default
//            if fileManager.fileExists(atPath: outputPath.path) {
//                try fileManager.removeItem(at: outputPath)
//            }
//            
//            // 创建目标文件
//            fileManager.createFile(atPath: outputPath.path, contents: nil)
//            
//            // 按顺序合并所有分片
//            let fileHandle = try FileHandle(forWritingTo: outputPath)
//            defer { fileHandle.closeFile() }
//            
//            // 按顺序合并片段
//            for segmentPath in segmentPaths.sorted(by: { $0.lastPathComponent < $1.lastPathComponent }) {
//                let segmentData = try Data(contentsOf: segmentPath)
//                outputHandle.write(segmentData)
//                
//                // 清理临时文件
//                try fileManager.removeItem(at: segmentPath)
//            }
//            
//            task.completion(.success(outputPath))
//            task.state = .completed
//            
//        } catch {
//            task.completion(.failure(error))
//        }
//    }
//}

// M3U8处理器（支持加密）
class M3U8Handler: ProtocolHandler {
    var task: DownloadTask
    private var segmentTasks: [URLSessionDownloadTask] = []
    private var session: URLSession!
    private var parser = M3U8Parser()
    private var decryptor = AESDecryptor()
    private var keyManager = KeyManager.shared
    private var isCancelled = false
    
    init(task: DownloadTask) {
        self.task = task
        self.session = URLSession(
            configuration: .ephemeral,
            delegate: nil,
            delegateQueue: .global(qos: .userInitiated)
    }
    
    func startDownload(manager: DownloadManager) {
        isCancelled = false
        manager.updateTaskState(id: task.id, state: .downloading(progress: 0))
        
        // 下载并解析m3u8索引
        let indexTask = session.dataTask(with: task.url) { [weak self] data, response, error in
            guard let self = self, let data = data, let content = String(data: data, encoding: .utf8) else {
                manager.updateTaskState(id: self?.task.id ?? UUID(), state: .failed(error: error?.localizedDescription ?? "Invalid M3U8 content"))
                return
            }
            
            let baseURL = task.url.deletingLastPathComponent()
            let encryptedM3U8 = self.parser.parseEncryptedM3U8(content: content, baseURL: baseURL)
            
            if encryptedM3U8.method == "AES-128" {
                self.downloadEncryptedSegments(manager: manager, encryptedM3U8: encryptedM3U8)
            } else {
                self.downloadSegments(manager: manager, segments: encryptedM3U8.tsSegments)
            }
        }
        indexTask.resume()
    }
    
    private func downloadEncryptedSegments(manager: DownloadManager, encryptedM3U8: EncryptedM3U8) {
        guard let keyURI = encryptedM3U8.keyURI else {
            manager.updateTaskState(id: task.id, state: .failed(error: "Missing encryption key"))
            return
        }
        
        keyManager.getKey(for: keyURI) { [weak self] result in
            switch result {
            case .success(let keyData):
                let iv = encryptedM3U8.iv ?? self?.generateDefaultIV(for: keyData) ?? Data()
                self?.downloadAndDecryptSegments(
                    manager: manager,
                    segments: encryptedM3U8.tsSegments,
                    key: keyData,
                    iv: iv
                )
                
            case .failure(let error):
                manager.updateTaskState(id: self?.task.id ?? UUID(), state: .failed(error: error.localizedDescription))
            }
        }
    }
    
    private func downloadAndDecryptSegments(manager: DownloadManager, segments: [URL], key: Data, iv: Data) {
        let group = DispatchGroup()
        let tempDir = manager.fileManager.createTempDirectory(for: task)
        var segmentPaths: [URL] = []
        var errors: [Error] = []
        
        for (index, segmentURL) in segments.enumerated() {
            if isCancelled { break }
            
            group.enter()
            let segmentIV = generateSegmentIV(baseIV: iv, sequence: index)
            
            let segmentTask = session.downloadTask(with: segmentURL) { [weak self] location, response, error in
                defer { group.leave() }
                
                guard let self = self else { return }
                if self.isCancelled { return }
                
                do {
                    guard let location = location else {
                        throw DownloadError.invalidSegment
                    }
                    
                    let segmentData = try Data(contentsOf: location)
                    let decryptedData = try self.decryptor.decryptTS(data: segmentData, key: key, iv: segmentIV)
                    
                    guard self.decryptor.validateDecryption(decryptedData: decryptedData) else {
                        throw DownloadError.decryptionFailed
                    }
                    
                    let segmentPath = tempDir.appendingPathComponent("seg_\(index).ts")
                    try decryptedData.write(to: segmentPath)
                    segmentPaths.append(segmentPath)
                    
                    DispatchQueue.main.async {
                        let progress = Double(index+1) / Double(segments.count)
                        manager.updateTaskState(id: self.task.id, state: .downloading(progress: progress * 0.9))
                    }
                } catch {
                    errors.append(error)
                }
            }
            segmentTask.resume()
            segmentTasks.append(segmentTask)
        }
        
        group.notify(queue: .global()) { [weak self] in
            guard let self = self else { return }
            
            if !errors.isEmpty || self.isCancelled {
                manager.updateTaskState(id: self.task.id, state: .failed(error: errors.first?.localizedDescription ?? "Download cancelled"))
                return
            }
            
            DispatchQueue.main.async {
                manager.updateTaskState(id: self.task.id, state: .merging)
            }
            
            self.mergeSegments(manager: manager, segmentPaths: segmentPaths)
        }
    }
    
    private func downloadSegments(manager: DownloadManager, segments: [URL]) {
        let group = DispatchGroup()
        let tempDir = manager.fileManager.createTempDirectory(for: task)
        var segmentPaths: [URL] = []
        var errors: [Error] = []
        
        for (index, segmentURL) in segments.enumerated() {
            if isCancelled { break }
            
            group.enter()
            let segmentTask = session.downloadTask(with: segmentURL) { [weak self] location, response, error in
                defer { group.leave() }
                
                guard let self = self else { return }
                if self.isCancelled { return }
                
                do {
                    guard let location = location else {
                        throw DownloadError.invalidSegment
                    }
                    
                    let segmentPath = tempDir.appendingPathComponent("seg_\(index).ts")
                    try FileManager.default.moveItem(at: location, to: segmentPath)
                    segmentPaths.append(segmentPath)
                    
                    DispatchQueue.main.async {
                        let progress = Double(index+1) / Double(segments.count)
                        manager.updateTaskState(id: self.task.id, state: .downloading(progress: progress * 0.9))
                    }
                } catch {
                    errors.append(error)
                }
            }
            segmentTask.resume()
            segmentTasks.append(segmentTask)
        }
        
        group.notify(queue: .global()) { [weak self] in
            guard let self = self else { return }
            
            if !errors.isEmpty || self.isCancelled {
                manager.updateTaskState(id: self.task.id, state: .failed(error: errors.first?.localizedDescription ?? "Download cancelled"))
                return
            }
            
            DispatchQueue.main.async {
                manager.updateTaskState(id: self.task.id, state: .merging)
            }
            
            self.mergeSegments(manager: manager, segmentPaths: segmentPaths)
        }
    }
    
    private func mergeSegments(manager: DownloadManager, segmentPaths: [URL]) {
        do {
            let sortedSegments = segmentPaths.sorted {
                guard let idx1 = Int($0.deletingPathExtension().lastPathComponent.split(separator: "_").last ?? ""),
                      let idx2 = Int($1.deletingPathExtension().lastPathComponent.split(separator: "_").last ?? "") else {
                    return false
                }
                return idx1 < idx2
            }
            
            let outputPath = manager.fileManager.localPath(for: task)
            if FileManager.default.fileExists(atPath: outputPath.path) {
                try FileManager.default.removeItem(at: outputPath)
            }
            
            FileManager.default.createFile(atPath: outputPath.path, contents: nil)
            let outputHandle = try FileHandle(forWritingTo: outputPath)
            
            for segmentPath in sortedSegments {
                let segmentData = try Data(contentsOf: segmentPath)
                outputHandle.write(segmentData)
                try FileManager.default.removeItem(at: segmentPath)
            }
            
            try outputHandle.close()
            try FileManager.default.removeItem(at: outputPath.deletingLastPathComponent())
            
            manager.updateTask(id: task.id, state: .completed, localPath: outputPath)
        } catch {
            manager.updateTaskState(id: task.id, state: .failed(error: error.localizedDescription))
        }
    }
    
    private func generateDefaultIV(for key: Data) -> Data {
        return key.md5()
    }
    
    private func generateSegmentIV(baseIV: Data, sequence: Int) -> Data {
        var mutableIV = baseIV
        for i in 0..<min(MemoryLayout<Int>.size, baseIV.count) {
            mutableIV[i] ^= UInt8((sequence >> (i * 8)) & 0xFF)
        }
        return mutableIV
    }
    
    func pauseDownload() {
        segmentTasks.forEach { $0.cancel() }
    }
    
    func resumeDownload() {
        // M3U8不支持真正的暂停/恢复，需要重新开始
        startDownload(manager: DownloadManager.shared)
    }
    
    func cancelDownload() {
        isCancelled = true
        segmentTasks.forEach { $0.cancel() }
    }
}
